package org.example;


import org.pcap4j.packet.IpPacket;
import org.pcap4j.packet.IpSelector;
import org.pcap4j.packet.UdpPacket;

import java.net.Inet4Address;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.UnknownHostException;
import java.nio.ByteBuffer;
import java.nio.channels.AsynchronousServerSocketChannel;
import java.nio.channels.AsynchronousSocketChannel;
import java.util.Optional;
import java.util.concurrent.*;

public class ServerMain {

    private static InetAddress inetAddress = null;

    private static ConcurrentHashMap<Integer,AsynchronousSocketChannel> channelMap = new ConcurrentHashMap<>();


    // 创建线程池，使用虚拟线程
    private static final ExecutorService executorService = Executors.newVirtualThreadPerTaskExecutor();

    static {
        try {
            String address = Optional.ofNullable(System.getenv("ADDRESS")).orElse("172.29.1.1");
            inetAddress = InetAddress.getByName(address);
        } catch (UnknownHostException e) {
            e.printStackTrace();
        }
    }


    public static void main(String[] args) {
        // 设置端口号
        int port = Integer.parseInt(Optional.ofNullable(System.getenv("PORT")).orElse("9999"));
        // 创建异步服务器套接字通道
        try (AsynchronousServerSocketChannel serverChannel = AsynchronousServerSocketChannel.open()) {
            // 绑定到本地端口
            serverChannel.bind(new InetSocketAddress(port));

            System.out.println("服务器启动，监听端口 " + port);

            // 接受连接
            while (true) {
                // 接受客户端连接
                AsynchronousSocketChannel clientChannel = serverChannel.accept().get();
                System.out.println("新连接： " + clientChannel.getRemoteAddress());

                // 使用虚拟线程处理每个连接
                executorService.submit(() -> handleClient(clientChannel));
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static void handleClient(AsynchronousSocketChannel clientChannel) {
        try {

            pushChannelMap(clientChannel);

            var buff = ByteBuffer.allocate(4);
            while (!Thread.interrupted()){
                var len = clientChannel.read(buff).get();
                if(len <= 0){
                    break;
                }
                buff.flip();
                var anInt = buff.getInt();
                var buff_dy = ByteBuffer.allocate(anInt);
                var len_dy = clientChannel.read(buff_dy).get();
                if(len_dy <= 0){
                    break;
                }
                var bt = new byte[len_dy];
                buff_dy.flip();
                buff_dy.get(bt);
                parsePacket(bt);
                buff.clear();
            }

        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                clientChannel.close();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    private static void pushChannelMap(AsynchronousSocketChannel clientChannel) throws Exception {
        var address = inetAddress.getAddress();
        var wrap = ByteBuffer.wrap(address);
        var anInt = wrap.getInt();

        for (int i = 0; i < 255; i++) {
            anInt++;

            if(channelMap.get(anInt)==null){

                channelMap.put(anInt,clientChannel);

                var bb = ByteBuffer.allocate(4);
                bb.putInt(anInt);
                bb.flip();

                clientChannel.write(bb).get();
                return;
            }

        }
        throw new RuntimeException("找不到可用的IP");

    }

    private static void parsePacket(byte[] data){

        try {
            IpPacket packet = (IpPacket) IpSelector.newPacket(data, 0, data.length);
            InetAddress dstAddr = packet.getHeader().getDstAddr();
            InetAddress srcAddr = packet.getHeader().getSrcAddr();

            if(dstAddr instanceof Inet4Address){
                Inet4Address dst4Addr = (Inet4Address)dstAddr;
                Inet4Address src4Addr = (Inet4Address)srcAddr;

                ByteBuffer srcwrap = ByteBuffer.wrap(src4Addr.getAddress());
                var srcKey = srcwrap.getInt();

                ByteBuffer allocate = ByteBuffer.allocate(data.length + 4);
                allocate.putInt(data.length);
                allocate.put(data);
                allocate.flip();

                if(dst4Addr.getAddress()[3] == (byte) 255){
                    //广播发送给全部channel
                    for (var entry : channelMap.entrySet()) {
                        var key = entry.getKey();
                        var value = entry.getValue();

                        if(key == srcKey){
                            //不需要广播自己
                            continue;
                        }

                        try {
                            value.write(allocate).get();
                        }catch (Exception e){
                            e.printStackTrace();
                            value.close();
                            channelMap.remove(key);
                        }

                    }

                }else{
                    ByteBuffer wrap = ByteBuffer.wrap(dst4Addr.getAddress());
                    var Key = wrap.getInt();
                    var asc = channelMap.get(Key);
                    if(asc!=null){
                        try {
                            asc.write(allocate).get();
                        }catch (Exception e){
                            e.printStackTrace();
                            asc.close();
                            channelMap.remove(Key);
                        }
                    }

                }
            }


            if(packet.getPayload() instanceof UdpPacket){

            }


        }catch (Exception e){
            e.printStackTrace();
        }
    }
}